# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

import numpy as np

import torch
import torch.nn as nn
from sympy.physics.units import velocity
from torch.distributions import Normal
from torch.nn.modules import rnn
from rsl_rl.modules.him_estimator import HIMEstimator
from rsl_rl.utils import unpad_trajectories

class ActorCritic(nn.Module):
    is_recurrent = False
    contain_student = True
    def __init__(self,  num_actor_obs,
                        num_critic_obs,
                        num_actions,
                        num_STMEnc,
                        num_state_estimator,
                        num_estimator,
                        num_terrain_obs,
                        num_privilege,
                        command_index,
                        student_train = False,
                        actor_hidden_dims=[256, 256, 256],
                        critic_hidden_dims=[256, 256, 256],
                        terrain_hidden_dims=[256, 128],
                        privilege_hidden_dims=[64, 32],
                        STMEnc_hidden_dims=[256, 128],
                        state_estimator_hidden_dims = [64, 32],
                        heights_encoder_hidden_dims = [256, 128],
                        lstm_encoder_hidden_dims = [256, 128],
                        hidden_size=0,
                        lstm_num_layers=0,
                        num_out_terrain = 0,
                        num_out_privilege = 0,
                        num_obs = 45,
                        activation='elu',
                        init_noise_std=1.0,
                        fixed_std=False,
                        **kwargs):
        if kwargs:
            print("ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs.keys()]))
        super(ActorCritic, self).__init__()

        self.student_train = student_train
        activation = get_activation(activation)

        # # Policy
        if num_STMEnc is not None:
            STMEnc_layers = []
            STMEnc_layers.append(nn.Linear(num_STMEnc, STMEnc_hidden_dims[0]))
            STMEnc_layers.append(activation)
            for l in range(len(STMEnc_hidden_dims)):
                if l == len(STMEnc_hidden_dims) - 1:
                    STMEnc_layers.append(nn.Linear(STMEnc_hidden_dims[l], 32))
                else:
                    STMEnc_layers.append(nn.Linear(STMEnc_hidden_dims[l], STMEnc_hidden_dims[l + 1]))
                    STMEnc_layers.append(activation)
            self.STMEnc = nn.Sequential(*STMEnc_layers)
            print(f"STMEnc MLP: {self.STMEnc}")

        if num_state_estimator is not None:
            self.command_index = command_index
            state_estimator_layers = []
            state_estimator_layers.append(nn.Linear(num_state_estimator, state_estimator_hidden_dims[0]))
            state_estimator_layers.append(activation)
            for l in range(len(state_estimator_hidden_dims)):
                if l == len(state_estimator_hidden_dims) - 1:
                    state_estimator_layers.append(nn.Linear(state_estimator_hidden_dims[l], 3))
                else:
                    state_estimator_layers.append(nn.Linear(state_estimator_hidden_dims[l], state_estimator_hidden_dims[l + 1]))
                    state_estimator_layers.append(activation)
            self.state_estimator = nn.Sequential(*state_estimator_layers)
            print(f"state_estimator MLP: {self.state_estimator}")

        # Estimator
        if num_estimator is not None:
            self.estimator = HIMEstimator(temporal_steps=5, num_one_step_obs=45)
            print(f'Estimator: {self.estimator.encoder}')

        if num_terrain_obs is not None:
            # Policy privilege
            terrain_layers = []
            terrain_layers.append(nn.Linear(num_terrain_obs, terrain_hidden_dims[0]))
            terrain_layers.append(activation)
            for l in range(len(terrain_hidden_dims)):
                if l == len(terrain_hidden_dims) - 1:
                    terrain_layers.append(nn.Linear(terrain_hidden_dims[l], num_out_terrain))
                else:
                    terrain_layers.append(nn.Linear(terrain_hidden_dims[l], terrain_hidden_dims[l + 1]))
                    terrain_layers.append(activation)
            self.terrain = nn.Sequential(*terrain_layers)
            print(f'terrain: {self.terrain}')

            self.lstm = Memory(num_obs, lstm_num_layers, hidden_size)
            self.fc = nn.Linear(hidden_size, num_out_terrain + num_out_privilege)
            self.criterion = nn.MSELoss()
            self.num_out_terrain = num_out_terrain
            print(f"LSTM: {self.lstm}")
            print(f"LSTM_fc: {self.fc}")
        if num_privilege is not None:
            privilege_layers = []
            privilege_layers.append(nn.Linear(num_privilege, privilege_hidden_dims[0]))
            privilege_layers.append(activation)
            for l in range(len(privilege_hidden_dims)):
                if l == len(privilege_hidden_dims) - 1:
                    privilege_layers.append(nn.Linear(privilege_hidden_dims[l], num_out_privilege))
                else:
                    privilege_layers.append(nn.Linear(privilege_hidden_dims[l], privilege_hidden_dims[l + 1]))
                    privilege_layers.append(activation)
            self.privilege = nn.Sequential(*privilege_layers)
            print(f'privilege: {self.privilege}')
            self.num_out_privilege = num_out_privilege

        mlp_input_dim_a = num_actor_obs
        mlp_input_dim_c = num_critic_obs

        # Policy
        actor_layers = []
        actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
        actor_layers.append(activation)
        for l in range(len(actor_hidden_dims)):
            if l == len(actor_hidden_dims) - 1:
                actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
            else:
                actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
                actor_layers.append(activation)
        self.actor = nn.Sequential(*actor_layers)

        # Value function
        critic_layers = []
        critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0]))
        critic_layers.append(activation)
        for l in range(len(critic_hidden_dims)):
            if l == len(critic_hidden_dims) - 1:
                critic_layers.append(nn.Linear(critic_hidden_dims[l], 1))
            else:
                critic_layers.append(nn.Linear(critic_hidden_dims[l], critic_hidden_dims[l + 1]))
                critic_layers.append(activation)
        self.critic = nn.Sequential(*critic_layers)

        print(f"Actor MLP: {self.actor}")
        print(f"Critic MLP: {self.critic}")

        # Action noise
        self.fixed_std = fixed_std
        std = init_noise_std * torch.ones(num_actions)
        self.std = torch.tensor(std) if fixed_std else nn.Parameter(std)
        self.distribution = None
        # disable args validation for speedup
        Normal.set_default_validate_args = False
        
        # seems that we get better performance without init
        # self.init_memory_weights(self.memory_a, 0.001, 0.)
        # self.init_memory_weights(self.memory_c, 0.001, 0.)
        self.velocity = None
        self.heights_encode = None

    @staticmethod
    # not used at the moment
    def init_weights(sequential, scales):
        [torch.nn.init.orthogonal_(module.weight, gain=scales[idx]) for idx, module in
         enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))]


    def reset(self, dones=None):
        # self.memory_t.reset(dones)
        pass

    def forward(self):
        raise NotImplementedError
    
    @property
    def action_mean(self):
        return self.distribution.mean

    @property
    def action_std(self):
        return self.distribution.stddev
    
    @property
    def entropy(self):
        return self.distribution.entropy().sum(dim=-1)

    def update_distribution(self, observations, terrain_obs, privilege_obs):
        if hasattr(self, 'STMEnc'):
            command = observations[:, self.command_index:self.command_index + 3]
            history_encode = self.STMEnc(observations[:, 0:225])
            self.velocity = self.state_estimator(history_encode.detach())
            self.heights_encode = self.heights_encoder(observations[:, 225:])
            mean = self.actor(torch.cat([command, history_encode,self.velocity.detach(),self.heights_encode], dim=-1))
        if hasattr(self, 'estimator'):
            with torch.no_grad():
                latent = self.estimator(observations[:, 0:225])
            terrain_latent = self.terrain(terrain_obs)
            privilege_latent = self.privilege(privilege_obs)
            mean = self.actor(torch.cat((observations, latent, terrain_latent, privilege_latent), dim=-1))
        else:
            mean = self.actor(observations)
        std = self.std.to(mean.device)
        self.distribution = Normal(mean, mean*0. + std)

    def update_student_distribution(self, observations, masks=None, hidden_states=None):
        if hidden_states is not None:
            lstm_out = self.lstm(observations, masks, hidden_states).squeeze(0)
            output = self.fc(lstm_out).reshape(-1, self.num_out_privilege + self.num_out_terrain)
        else:
            lstm_out = self.lstm(observations).squeeze(0)
            output = self.fc(lstm_out)
        if masks is not None:
            observations = unpad_trajectories(observations,masks).reshape(-1,observations.shape[-1])
        if hasattr(self, 'estimator'):
            with torch.no_grad():
                him_latent = self.estimator(observations)
        terrain_latent = output[:, :self.num_out_terrain]
        privilege_latent = output[:, -self.num_out_privilege:]
        mean = self.actor(torch.cat((observations, him_latent, terrain_latent, privilege_latent), dim=-1))
        std = self.std.to(mean.device)
        self.distribution = Normal(mean, mean * 0. + std)
        return output

    def act(self, obs, terrain_obs, privilege_obs, masks=None, hidden_states=None):
        if self.student_train is False:
            observations = obs
            terrain_obs = terrain_obs
            privilege_obs = privilege_obs
            self.update_distribution(observations, terrain_obs, privilege_obs)
        else:
            self.update_student_distribution(obs, masks, hidden_states)
        return self.distribution.sample()

    def student_act(self, obs):
        latent = self.update_student_distribution(obs)
        return latent, self.distribution.sample()

    def get_actions_log_prob(self, actions):
        return self.distribution.log_prob(actions).sum(dim=-1)

    def act_inference(self, observations,terrain_obs,privilege_obs):
        if hasattr(self, 'STMEnc'):
            command = observations[:, self.command_index:self.command_index + 3]
            history_encode = self.STMEnc(observations[:, 0:225])
            self.velocity = self.state_estimator(history_encode.detach())
            # print(self.velocity)
            self.heights_encode = self.heights_encoder(observations[:, 225:])
            actions_mean = self.actor(torch.cat([command, history_encode,self.velocity.detach(),self.heights_encode], dim=-1))
        if hasattr(self, 'estimator'):
            with torch.no_grad():
                latent = self.estimator(observations[:, 0:225])
            terrain_latent = self.terrain(terrain_obs)
            privilege_latent = self.privilege(privilege_obs)
            actions_mean = self.actor(torch.cat((observations, latent, terrain_latent, privilege_latent), dim=-1))
        else:
            actions_mean = self.actor(observations)
        return actions_mean

    def student_act_inference(self, observations):
        lstm_out = self.lstm(observations).squeeze(0)
        output = self.fc(lstm_out)
        him_latent = self.estimator(observations)
        terrain_latent = output[:, :self.num_out_terrain]
        privilege_latent = output[:, -self.num_out_privilege:]
        actions_mean = self.actor(torch.cat((observations, him_latent, terrain_latent, privilege_latent), dim=-1))
        return actions_mean

    def evaluate(self, observations, terrain_obs, privilege_obs, masks=None, hidden_states=None):
        # if self.student_train is False:
        if hasattr(self, 'estimator'):
            with torch.no_grad():
                him_latent = self.estimator(observations)
        terrain_latent = self.terrain(terrain_obs)
        privilege_latent = self.privilege(privilege_obs)
        value = self.critic(torch.cat((observations, him_latent, terrain_latent, privilege_latent), dim=-1))
        # else:
        #     if hidden_states is not None:
        #         lstm_out = self.lstm(observations, masks, hidden_states).squeeze(0)
        #         output = self.fc(lstm_out).reshape(-1, self.num_out_privilege + self.num_out_terrain)
        #     else:
        #         lstm_out = self.lstm(observations).squeeze(0)
        #         output = self.fc(lstm_out)
        #     if masks is not None:
        #         observations = unpad_trajectories(observations,masks).reshape(-1,observations.shape[-1])
        #     him_latent = self.estimator(observations)
        #     terrain_latent = output[:, :self.num_out_terrain]
        #     privilege_latent = output[:, -self.num_out_privilege:]
        #     critic_obs = torch.cat((observations, him_latent, terrain_latent, privilege_latent), dim=-1)
        #     value = self.critic(critic_obs)
        return value

    def get_hidden_states(self):
        return self.lstm.hidden_states

    def student_latent(self, observations, masks=None, hidden_states=None):
        lstm_out = self.lstm(observations, masks, hidden_states)
        output = self.fc(lstm_out)
        return output

    def teacher_latent(self, obs_terrain, privilege_obs):
        terrain_latent = self.terrain(obs_terrain)
        privilege_latent = self.privilege(privilege_obs)
        return torch.cat((terrain_latent, privilege_latent), dim=-1)

def get_activation(act_name):
    if act_name == "elu":
        return nn.ELU()
    elif act_name == "selu":
        return nn.SELU()
    elif act_name == "relu":
        return nn.ReLU()
    elif act_name == "crelu":
        return nn.ReLU()
    elif act_name == "lrelu":
        return nn.LeakyReLU()
    elif act_name == "tanh":
        return nn.Tanh()
    elif act_name == "sigmoid":
        return nn.Sigmoid()
    else:
        print("invalid activation function!")
        return None

class Memory(torch.nn.Module):
    def __init__(self, input_size, num_layers=1, hidden_size=256):
        super().__init__()
        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)
        self.hidden_states = None
        self.num_layers = num_layers
        self.hidden_size = hidden_size

    def forward(self, input, masks=None, hidden_states=None):
        batch_mode = masks is not None
        if batch_mode:
            # batch mode (policy update): need saved hidden states
            if hidden_states is None:
                raise ValueError("Hidden states not passed to memory module during policy update")
            out, _ = self.lstm(input, hidden_states)
            out = unpad_trajectories(out, masks)
        else:
            # inference mode (collection): use hidden states of last step
            out, self.hidden_states = self.lstm(input.unsqueeze(0), self.hidden_states)
        return out

    def reset(self, dones=None):
        # When the RNN is an LSTM, self.hidden_states_a is a list with hidden_state and cell_state
        for hidden_state in self.hidden_states:
            hidden_state[..., dones == 1, :] = 0.0